/**
 * @file layerEarth_MT.cpp
 * @author Yi Zhang (zhangyiss@icloud.com)
 * @brief A closed-form for plane wave electromagnetic fields (PWEF) on a n-layers Earth model with an oblique angle at a given frequency
 * @version 1.0
 * @date 2021-11-30
 * 
 * @copyright Copyright (c) 2021
 * 
 * +z downward, z=0 the Air-Earth interface
 * time dependence iwt
 * TE-mode: E = (Ex,0 ,0 ), H = (0, Hy,Hz)
 * TM-mode: E = (0 ,Ey,Ez), H = (Hx,0 ,0 )
 *          x (Ex for TE mode and Hx for TM mode)
 *        /
 *      /
 *    /    air (layer 0)
 *  /
 * O---------------------> y
 * |   layer 1
 * |---------------------
 * |   layer 2
 * |---------------------
 * |   ...
 * |---------------------
 * |   layer n
 * |
 * z
 * 
 */

#include "layerEarth_MT.h"

#include "cmath"
#include "exception"
#include "stdexcept"
#include "iostream"

#define pi 3.1415926535897932384626433832795
#define mu0 (4e-7*pi)
#define epsilon0 8.8541878176e-12

const static std::complex<double> _1i = std::complex<double>(0.0, 1.0);
const static std::complex<double> _zero = std::complex<double>(0.0, 0.0);
const static std::complex<double> _one = std::complex<double>(1.0, 0.0);

void matvec(const matrix2x2 &mat, const std::complex<double> &U, const std::complex<double> &D, 
    std::complex<double> &U_out, std::complex<double> &D_out)
{
    U_out = mat.val[0][0]*U + mat.val[0][1]*D;
    D_out = mat.val[1][0]*U + mat.val[1][1]*D;
    return;
}

void matmat(const matrix2x2 &mat1, const matrix2x2 &mat2, matrix2x2 &mat_out)
{
    mat_out.val[0][0] = mat1.val[0][0]*mat2.val[0][0] + mat1.val[0][1]*mat2.val[1][0];
    mat_out.val[0][1] = mat1.val[0][0]*mat2.val[0][1] + mat1.val[0][1]*mat2.val[1][1];
    mat_out.val[1][0] = mat1.val[1][0]*mat2.val[0][0] + mat1.val[1][1]*mat2.val[1][0];
    mat_out.val[1][1] = mat1.val[1][0]*mat2.val[0][1] + mat1.val[1][1]*mat2.val[1][1];
    return;
}

conductive_layer::conductive_layer(double c, double e, double m, double d)
{
    set(c, e, m, d);
}

void conductive_layer::set(double c, double e, double m, double d)
{
    cndt = c; epsl = e; mu = m; btm_depth = d;
}

layerEarthMT::layerEarthMT()
{
    initialized_ = false;
    omega_ = 2.0*pi;
    theta_ = 0.0;
    TM_mode_ = false;
    B0_ = 1.0;

    ky_ = Ex_ = Ey_ = Ez_ = Hx_ = Hy_ = Hz_ = _zero;
    Ex_y_ = Ex_z_ = Hx_y_ = Hx_z_ = _zero;
}

layerEarthMT::layerEarthMT(const std::vector<conductive_layer> &layers) : layerEarthMT()
{
    set_layers(layers);
}

void layerEarthMT::initialize()
{
    // assemble all layers
    int layer_size = all_layers_.size() - 1;

    if (all_layers_[0].btm_depth >= all_layers_[1].btm_depth)
    {
        throw std::runtime_error("Invalid air layer's bottom depth.");
        return;
    }

    if (all_layers_[layer_size].btm_depth <= all_layers_[layer_size-1].btm_depth)
    {
        throw std::runtime_error("Invalid bottom layer's bottom depth.");
        return;
    }

    // calculate wave numbers and thicknesses
    if (k_.size() != layer_size+1)
    {
        k_.clear(); k_.resize(layer_size+1);
        kz_.clear(); kz_.resize(layer_size+1);
        h_.clear(); h_.resize(layer_size+1);
        zhat_.clear(); zhat_.resize(layer_size+1);
        yhat_.clear(); yhat_.resize(layer_size+1);
        Z_.clear(); Z_.resize(layer_size+1);
        Y_.clear(); Y_.resize(layer_size+1);
        U_.clear(); U_.resize(layer_size+1);
        D_.clear(); D_.resize(layer_size+1);
        T_.clear(); T_.resize(layer_size+1);
    }

    for (int i = 0; i <= layer_size; i++)
    {
        zhat_[i] = -1.0*_1i*omega_*all_layers_[i].mu;
        yhat_[i] = all_layers_[i].cndt - _1i*omega_*all_layers_[i].epsl;
        k_[i] = std::sqrt(-1.0*zhat_[i]*yhat_[i]);
    }

    std::complex<double> k_air = k_[0];
    ky_ = k_air*std::sin(theta_);
    for (int i = 0; i <= layer_size; i++)
    {
        kz_[i] = std::sqrt(k_[i]*k_[i] - ky_*ky_);
        Z_[i] = -1.0*_1i*kz_[i]/yhat_[i];
        Y_[i] = _1i*kz_[i]/zhat_[i];
    }

    h_[0] = h_[layer_size] = 1e+50; // 这两个值其实没用到
    for (int i = 1; i < layer_size; i++)
    {
        if (all_layers_[i].btm_depth <= all_layers_[i-1].btm_depth)
        {
            throw std::runtime_error("Invalid layer depths.");
        }

        h_[i] = all_layers_[i].btm_depth - all_layers_[i-1].btm_depth;
    }

    for (int i = 1; i <= layer_size-1; i++)
    {
        if (TM_mode_)
        {
            T_[i].val[0][0] = 0.5*std::exp(     _1i*kz_[i]*h_[i])*(1.0 + Z_[i]/Z_[i-1]);
            T_[i].val[0][1] = 0.5*std::exp(-1.0*_1i*kz_[i]*h_[i])*(1.0 - Z_[i]/Z_[i-1]);
            T_[i].val[1][0] = 0.5*std::exp(     _1i*kz_[i]*h_[i])*(1.0 - Z_[i]/Z_[i-1]);
            T_[i].val[1][1] = 0.5*std::exp(-1.0*_1i*kz_[i]*h_[i])*(1.0 + Z_[i]/Z_[i-1]);
        }
        else // TE mode
        {
            T_[i].val[0][0] = 0.5*std::exp(     _1i*kz_[i]*h_[i])*(1.0 + Y_[i]/Y_[i-1]);
            T_[i].val[0][1] = 0.5*std::exp(-1.0*_1i*kz_[i]*h_[i])*(1.0 - Y_[i]/Y_[i-1]);
            T_[i].val[1][0] = 0.5*std::exp(     _1i*kz_[i]*h_[i])*(1.0 - Y_[i]/Y_[i-1]);
            T_[i].val[1][1] = 0.5*std::exp(-1.0*_1i*kz_[i]*h_[i])*(1.0 + Y_[i]/Y_[i-1]);
        }
    }

    int t = layer_size;
    if (TM_mode_)
    {
        T_[t].val[0][0] = 0.5*std::exp(-1.0*_1i*kz_[t]*all_layers_[t-1].btm_depth)*(1.0 + Z_[t]/Z_[t-1]);
        T_[t].val[0][1] = 0.5*std::exp(     _1i*kz_[t]*all_layers_[t-1].btm_depth)*(1.0 - Z_[t]/Z_[t-1]);
        T_[t].val[1][0] = 0.5*std::exp(-1.0*_1i*kz_[t]*all_layers_[t-1].btm_depth)*(1.0 - Z_[t]/Z_[t-1]);
        T_[t].val[1][1] = 0.5*std::exp(     _1i*kz_[t]*all_layers_[t-1].btm_depth)*(1.0 + Z_[t]/Z_[t-1]);
    }
    else // TE mode
    {
        T_[t].val[0][0] = 0.5*std::exp(-1.0*_1i*kz_[t]*all_layers_[t-1].btm_depth)*(1.0 + Y_[t]/Y_[t-1]);
        T_[t].val[0][1] = 0.5*std::exp(     _1i*kz_[t]*all_layers_[t-1].btm_depth)*(1.0 - Y_[t]/Y_[t-1]);
        T_[t].val[1][0] = 0.5*std::exp(-1.0*_1i*kz_[t]*all_layers_[t-1].btm_depth)*(1.0 - Y_[t]/Y_[t-1]);
        T_[t].val[1][1] = 0.5*std::exp(     _1i*kz_[t]*all_layers_[t-1].btm_depth)*(1.0 + Y_[t]/Y_[t-1]);
    }

    matrix2x2 S_out, S = T_[1];
    for (int i = 2; i <= layer_size; i++)
    {
        matmat(S, T_[i], S_out);
        S = S_out;
    }

    std::complex<double> S22 = S.val[1][1], S12 = S.val[0][1];
    U_[0] = S12*B0_/S22;
    D_[0] = B0_;
    U_[layer_size] = _zero;
    D_[layer_size] = B0_*_one/S22;

    std::complex<double> u, d, u_out, d_out;
    for (int i = layer_size-1; i >= 1; i--)
    {
        u = U_[i+1];
        d = D_[i+1];
        matvec(T_[i+1], u, d, u_out, d_out);
        U_[i] = u_out;
        D_[i] = d_out;
    }

    initialized_ = true;
    return;
}

void layerEarthMT::calculate_EH(double z, double y)
{
    if (!initialized_)
    {
        throw std::runtime_error("Current model setup is not initialized. call the function initialize() firstly.");
        return;
    }

    int layer_size = all_layers_.size() - 1;

    // determine the layer that z is located inside
    int hl = 0; // host layer
    while (z > all_layers_[hl].btm_depth)
    {
        if (hl == layer_size) break;
        hl++;
    }

    if (TM_mode_)
    {
        if (hl != layer_size)
        {
            Hx_ = std::exp(_1i*ky_*y)*(U_[hl]*std::exp(-1.0*_1i*kz_[hl]*(z - all_layers_[hl].btm_depth)) 
                + D_[hl]*std::exp(_1i*kz_[hl]*(z - all_layers_[hl].btm_depth)));
            // Hy_ = _zero;
            // Hz_ = _zero;

            // Ex_ = _zero;
            Ey_ = std::exp(_1i*ky_*y)*Z_[hl]*(U_[hl]*std::exp(-1.0*_1i*kz_[hl]*(z - all_layers_[hl].btm_depth)) 
                - D_[hl]*std::exp(_1i*kz_[hl]*(z - all_layers_[hl].btm_depth)));
            Ez_ = Hx_*(-1.0*_1i*ky_)/yhat_[hl];
        }
        else
        {
            Hx_ = std::exp(_1i*ky_*y)*(U_[hl]*std::exp(-1.0*_1i*kz_[hl]*z) + D_[hl]*std::exp(_1i*kz_[hl]*z));
            // Hy_ = _zero;
            // Hz_ = _zero;

            // Ex_ = _zero;
            Ey_ = std::exp(_1i*ky_*y)*Z_[hl]*(U_[hl]*std::exp(-1.0*_1i*kz_[hl]*z) - D_[hl]*std::exp(_1i*kz_[hl]*z));
            Ez_ = Hx_*(-1.0*_1i*ky_)/yhat_[hl];
        }

        Hx_y_ = -1.0*yhat_[hl]*Ez_;
        Hx_z_ = yhat_[hl]*Ey_;
    }
    else // TE mode
    {
        if (hl != layer_size)
        {
            Ex_ = std::exp(_1i*ky_*y)*(U_[hl]*std::exp(-1.0*_1i*kz_[hl]*(z - all_layers_[hl].btm_depth)) 
                + D_[hl]*std::exp(_1i*kz_[hl]*(z - all_layers_[hl].btm_depth)));
            // Ey_ = _zero;
            // Ez_ = _zero;

            // Hx_ = _zero;
            Hy_ = std::exp(_1i*ky_*y)*Y_[hl]*(U_[hl]*std::exp(-1.0*_1i*kz_[hl]*(z - all_layers_[hl].btm_depth)) 
                - D_[hl]*std::exp(_1i*kz_[hl]*(z - all_layers_[hl].btm_depth)));
            Hz_ = Ex_*(_1i*ky_)/zhat_[0];
        }
        else
        {
            Ex_ = std::exp(_1i*ky_*y)*(U_[hl]*std::exp(-1.0*_1i*kz_[hl]*z) + D_[hl]*std::exp(_1i*kz_[hl]*z));
            // Ey_ = _zero;
            // Ez_ = _zero;

            // Hx_ = _zero;
            Hy_ = std::exp(_1i*ky_*y)*Y_[hl]*(U_[hl]*std::exp(-1.0*_1i*kz_[hl]*z) - D_[hl]*std::exp(_1i*kz_[hl]*z));
            Hz_ = Ex_*(_1i*ky_)/zhat_[0];
        }

        Ex_y_ = zhat_[hl]*Hz_;
        Ex_z_ = -1.0*zhat_[hl]*Hy_;
    }

    return;
}

std::complex<double> layerEarthMT::get_Ex(){return Ex_;}
std::complex<double> layerEarthMT::get_Ey(){return Ey_;}
std::complex<double> layerEarthMT::get_Ez(){return Ez_;}
std::complex<double> layerEarthMT::get_Hx(){return Hx_;}
std::complex<double> layerEarthMT::get_Hy(){return Hy_;}
std::complex<double> layerEarthMT::get_Hz(){return Hz_;}
std::complex<double> layerEarthMT::get_Ex_dy(){return Ex_y_;}
std::complex<double> layerEarthMT::get_Ex_dz(){return Ex_z_;}
std::complex<double> layerEarthMT::get_Hx_dy(){return Hx_y_;}
std::complex<double> layerEarthMT::get_Hx_dz(){return Hx_z_;}

double layerEarthMT::get_apparent_resistivity()
{
    if (TM_mode_) return std::norm(Ey_/Hx_)/(mu0*omega_);
    return std::norm(Ex_/Hy_)/(mu0*omega_);
}

double layerEarthMT::get_phase_angle()
{
    if (TM_mode_) return std::arg(Ey_/Hx_)*180.0/pi;
    return std::arg(Ex_/Hy_)*180.0/pi;
}

void layerEarthMT::set_layers(const std::vector<conductive_layer> &layers)
{
    if (layers.size() < 2)
    {
        throw std::runtime_error("There must be at least two layers. From layerEarthMT::set_layers(...)");
        return;
    }
    
    if (!all_layers_.empty()) all_layers_.clear();

    // add layers
    conductive_layer tmp_layer;
    for (int i = 0; i < layers.size(); i++)
    {
        if (layers[i].cndt <= 0.0 || layers[i].epsl <= 0.0 || layers[i].epsl <= 0.0)
        {
            throw std::invalid_argument("Invalid layer parameters.");
        }
        else
        {
            tmp_layer = layers[i];
            tmp_layer.mu *= mu0;
            tmp_layer.epsl *= epsilon0;
            all_layers_.push_back(tmp_layer);
        }
    }

    initialized_ = false;
    return;
}

void layerEarthMT::set_frequency(double frequency)
{
    if (frequency <= 0.0)
    {
        throw std::invalid_argument("Invalid frequency.");
    }

    omega_ = 2.0*pi*frequency;
    initialized_ = false;
    return;
}

void layerEarthMT::set_incident_angle(double theta)
{
    if (theta < 0.0 || theta > 90.0)
    {
        throw std::invalid_argument("Invalid incident angle.");
    }

    theta_ = theta;
    initialized_ = false;
    return;
}

void layerEarthMT::set_to_TM_mode()
{
    TM_mode_ = true;
    initialized_ = false;
    return;
}

void layerEarthMT::set_to_TE_mode()
{
    TM_mode_ = false;
    initialized_ = false;
    return;
}

void layerEarthMT::set_ground_amplitude(double b0)
{
    if (b0 <= 0.0)
    {
        throw std::invalid_argument("Invalid ground amplitude.");
    }

    B0_ = b0;
    initialized_ = false;
    return;
}
